package com.sicau.agent;

import cn.hutool.core.util.StrUtil;
import com.sicau.agent.model.AgentState;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;

@Data
@Slf4j
/**
 * 抽象基础代理类，用于管理代理状态和执行流程。
 * 提供状态转换、内存管理和基于步骤的执行循环的基础功能。
 * 子类必须实现step方法。
 */
public abstract class BaseAgent {
    /**
     * 聊天响应结果
     */
    private ChatResponse chatResponse;
    /**
     * 智能助手名称
     */
    private String name;
    /**
     * 系统提示语
     */
    private String systemPrompt;
    /**
     * 执行下一步操作的提示语
     */
    private String nextStepPrompt;
    /**
     * 聊天客户端
     */
    private ChatClient chatClient;
    /**
     * 消息上下文
     */
    private List<Message> messageList = new ArrayList<>();
    /**
     * 最大步骤数
     */
    private int maxStep = 10;
    /**
     * 智能助手状态
     */
    private AgentState agentState = AgentState.IDLE;

    /**
     * 定义智能助手的运行流程，AgentLoop
     * @param userPrompt 用户输入
     * @return 智能助手的输出
     */
    public String run(String userPrompt) {
        // 处理异常情况
        if (StrUtil.isBlank(userPrompt)) {
            log.error("用户输入为空！");
            return "请输入有效的用户提示！";
        }
        // 避免工具重复执行
        if (agentState != AgentState.IDLE) {
            log.error("智能体正在执行中！");
            return "智能体正在执行中！";
        }
        log.info("用户询问：{}", userPrompt);
        // 将智能助手状态设置为运行中
        agentState = AgentState.RUNNING;
        // 记录消息上下文
        messageList.add(new UserMessage(userPrompt));
        // 记录结果列表
        List<String> resultList = new ArrayList<>();
        int currentStep = 0;
        try {
            while (currentStep < maxStep && agentState != AgentState.FINISHED) {
                currentStep++;
                log.info("{}执行第{}步/{}", name, currentStep, maxStep);
                String result = step();
                log.info("{}第{}步执行结果：{}", name, currentStep, result);
                // 添加到结果列表
                resultList.add("第" + currentStep + "步执行结果：\n" + result);
            }
            if (currentStep >= maxStep) {
                log.info("{}超出最大执行步骤！", name);
                agentState = AgentState.FINISHED;
                resultList.add("超出最大执行步骤！");
            }
            return String.join("\n", resultList);
        } catch (Exception e) {
            log.error("智能助手执行异常！", e);
            agentState = AgentState.ERROR;
            return "智能助手执行异常！";
        } finally {
            clean();
        }
    }

    /**
     * 定义智能助手的运行流程，AgentLoop
     * @param userPrompt 用户输入
     * @return 智能助手的流式输出
     */
    public SseEmitter runWithSSE(String userPrompt) {
        SseEmitter sseEmitter = new SseEmitter(300000L);
        CompletableFuture.runAsync(() -> {
            try {
                // 处理异常情况
                if (StrUtil.isBlank(userPrompt)) {
                    log.error("用户输入为空！");
                    sseEmitter.send("请输入有效的用户提示！");
                    sseEmitter.complete();
                    return;
                }
                // 避免工具重复执行
                if (agentState != AgentState.IDLE) {
                    log.error("智能体正在执行中！");
                    sseEmitter.send("智能体正在执行中！");
                    sseEmitter.complete();
                    return;
                }
            } catch (IOException e) {
                log.error("发送数据异常！", e);
                sseEmitter.completeWithError(e);
            }
            log.info("用户询问：{}", userPrompt);
            // 将智能助手状态设置为运行中
            agentState = AgentState.RUNNING;
            // 记录消息上下文
            messageList.add(new UserMessage(userPrompt));
            int currentStep = 0;
            try {
                while (currentStep < maxStep && agentState != AgentState.FINISHED) {
                    currentStep++;
                    log.info("{}执行第{}步/{}", name, currentStep, maxStep);
                    sseEmitter.send(String.format("%s执行第%d步/%d", name, currentStep, maxStep));
                    String result = step();
                    log.info("{}第{}步执行结果：{}", name, currentStep, result);
                    sseEmitter.send(String.format("%s第%d步执行结果：\n%s", name, currentStep, result));
                }
                if (currentStep >= maxStep) {
                    log.info("{}超出最大执行步骤！", name);
                    agentState = AgentState.FINISHED;
                    sseEmitter.send(name + "超出最大执行步骤" + maxStep);
                }
            } catch (Exception e) {
                log.error("智能助手执行异常！", e);
                agentState = AgentState.ERROR;
                try {
                    sseEmitter.send(name + "执行异常！");
                } catch (IOException ex) {
                    sseEmitter.completeWithError(ex);
                }
            } finally {
                sseEmitter.complete();
                clean();
            }
        });
        sseEmitter.onTimeout(() -> {
            try {
                sseEmitter.send(name + "系统超时");
            } catch (IOException e) {
                sseEmitter.completeWithError(e);
            } finally {
                clean();
            }
        });
        return sseEmitter;
    }

    /**
     * 智能助手执行下一步操作
     * @return 单个步骤执行结果
     */
    protected abstract String step();

    /**
     * 清理数据
     */
    protected void clean() {
        log.info("清理数据");
    }
}
